{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 有点乱，大家根据自己习惯瞎改改吧\n",
    "- kernel替换完成后，不可逆，需要刷新重启\n",
    "- 测试结果都是取得稳定状态下的迭代速度，大概在20步之后"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# os.environ['CUDA_VISIBLE_DEVICES'] = '7'\n",
    "import torch\n",
    "import sys\n",
    "import importlib\n",
    "# 如果测huggingface中没有模型，需要把路径导进来，不然import了kernel，使用AutoModel加载模型也没有用处\n",
    "sys.path.append('/mnt/workspace/mdy/models')\n",
    "module = importlib.import_module('glm-4-9b-chat.modeling_chatglm')\n",
    "from transformers import AutoConfig, AutoModelForCausalLM, Qwen2ForCausalLM\n",
    "import time\n",
    "from copy import deepcopy\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# config = AutoConfig.from_pretrained('/mnt/workspace/mdy/models/Qwen2.5-7B-Instruct')\n",
    "dtype = torch.bfloat16\n",
    "bs = 8\n",
    "seq_len = 2048\n",
    "iters = 50\n",
    "model_path = '/mnt/workspace/mdy/models/glm-4-9b-chat'\n",
    "model_path = '/mnt/workspace/mdy/models/Qwen2.5-0.5B'\n",
    "# model_path = '/mnt/workspace/mdy/models/Llama-3.2-1B-Instruct'\n",
    "# model_path = '/mnt/workspace/mdy/models/Meta-Llama-3.1-8B-Instruct'\n",
    "model_path = '/mnt/workspace/mdy/models/glm-4-qwen2-format'\n",
    "model = AutoModelForCausalLM.from_pretrained(model_path,\n",
    "                                            _attn_implementation='flash_attention_2',\n",
    "                                           device_map='cuda', torch_dtype=dtype, trust_remote_code=True)\n",
    "# model = module.ChatGLMForConditionalGeneration.from_pretrained(model_path, _attn_implementation='flash_attention_2',\n",
    "#                                            device_map='cuda', torch_dtype=dtype, trust_remote_code=True)\n",
    "input_ids = torch.randint(10000, 20000, (bs, seq_len)).cuda()\n",
    "pos_ids = torch.arange(seq_len)[None, :].expand(bs, -1).contiguous().cuda()\n",
    "labels = input_ids\n",
    "grad_list = []\n",
    "for p in model.parameters():\n",
    "    grad_list.append(p)\n",
    "def zero_gard():\n",
    "    for p in grad_list:\n",
    "        if p.grad is not None:\n",
    "            p.grad.zero_()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50/50 [00:25<00:00,  1.93it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "25.961575031280518"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.empty_cache()\n",
    "t = time.time()\n",
    "for i in tqdm(range(iters)):\n",
    "    out = model(input_ids, position_ids=pos_ids, labels=labels)\n",
    "    out.loss.backward()\n",
    "t1 = time.time() - t\n",
    "t1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.0159,  0.0309, -0.0610,  ..., -0.0205,  0.0182, -0.0352],\n",
      "        [-0.0051,  0.0048, -0.0112,  ...,  0.0154,  0.0073, -0.0027],\n",
      "        [ 0.0123,  0.0124, -0.0742,  ..., -0.0075,  0.0182, -0.0182],\n",
      "        [-0.0019, -0.0076, -0.0201,  ...,  0.0031,  0.0033, -0.0094],\n",
      "        [-0.0094, -0.0095, -0.0092,  ..., -0.0120,  0.0060,  0.0026]],\n",
      "       device='cuda:0', dtype=torch.bfloat16)\n"
     ]
    }
   ],
   "source": [
    "print(grad_list[10].grad[:5])\n",
    "zero_gard()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mdy_triton.replace_kernel import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50/50 [00:21<00:00,  2.36it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "21.159708976745605"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.empty_cache()\n",
    "t = time.time()\n",
    "for i in tqdm(range(iters)):\n",
    "    out = model(input_ids, position_ids=pos_ids,labels=labels)\n",
    "    out.loss.backward()\n",
    "t2 = time.time() - t\n",
    "t2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.0203,  0.0364, -0.0806,  ..., -0.0184,  0.0064, -0.0302],\n",
      "        [-0.0008, -0.0011, -0.0292,  ...,  0.0112,  0.0091, -0.0053],\n",
      "        [ 0.0121,  0.0245, -0.0908,  ..., -0.0005,  0.0309, -0.0292],\n",
      "        [ 0.0025, -0.0101,  0.0013,  ...,  0.0003, -0.0092, -0.0123],\n",
      "        [-0.0061, -0.0008, -0.0120,  ..., -0.0122, -0.0010,  0.0017]],\n",
      "       device='cuda:0', dtype=torch.bfloat16)\n"
     ]
    }
   ],
   "source": [
    "print(grad_list[10].grad[:5])\n",
    "zero_gard()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# loss显存占用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "logits显存占用： 4.57763671875 G\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:21<00:00,  9.15it/s]\n"
     ]
    }
   ],
   "source": [
    "torch.cuda.empty_cache()\n",
    "bs = 8\n",
    "seq_len = 2048\n",
    "vocab_size = 150000\n",
    "dtype=torch.bfloat16\n",
    "factor = 4 if dtype == torch.float32 else 2\n",
    "print('logits显存占用：',(bs * seq_len * vocab_size) / (1024)**3 * factor,\"G\")\n",
    "logits = torch.randn(bs, seq_len, vocab_size, device='cuda', dtype=dtype, requires_grad=True)\n",
    "labels = torch.randint(0, 10, (bs, seq_len)).cuda()\n",
    "for i in tqdm(range(200)):\n",
    "    logits = logits.float()\n",
    "    shift_logits = logits[:, :-1].contiguous()\n",
    "    shift_labels = labels[:, 1:].contiguous()\n",
    "    loss_fn = torch.nn.CrossEntropyLoss()\n",
    "    loss = loss_fn(shift_logits.view(-1, vocab_size), shift_labels.view(-1))\n",
    "    loss.backward()\n",
    "    # break\n",
    "# 0.12s/it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "logits显存占用： 4.57763671875 G\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/200 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.6875,  0.2178,  0.1729,  ...,  0.6133, -2.2812,  0.4199],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<SelectBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/200 [00:01<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1.2460e-10, 3.0923e-10, 2.9468e-10,  ..., 4.5839e-10, 2.5352e-11,\n",
      "        3.7835e-10], device='cuda:0', dtype=torch.bfloat16,\n",
      "       grad_fn=<SelectBackward0>)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from mdy_triton.core import fast_cross_entropy_loss\n",
    "torch.cuda.empty_cache()\n",
    "bs = 8\n",
    "seq_len = 2048\n",
    "vocab_size = 150000\n",
    "dtype=torch.bfloat16\n",
    "factor = 4 if dtype == torch.float32 else 2\n",
    "print('logits显存占用：',(bs * seq_len * vocab_size) / (1024)**3 * factor,\"G\")\n",
    "extra_ignored_labels = torch.full((seq_len, 1), -100, device='cuda')\n",
    "logits = torch.randn(bs, seq_len, vocab_size, device='cuda', dtype=dtype, requires_grad=True)\n",
    "labels = torch.randint(0, vocab_size-1, (bs, seq_len)).cuda()\n",
    "for i in tqdm(range(200)):\n",
    "    shift_logits = logits\n",
    "    print(shift_logits[0][0])\n",
    "    shift_labels = torch.hstack((labels[..., 1:], extra_ignored_labels[:labels.shape[0]]))\n",
    "    loss = fast_cross_entropy_loss(shift_logits, shift_labels)\n",
    "    loss.backward()\n",
    "    print(shift_logits[0][0])\n",
    "    break\n",
    "# 重启运行\n",
    "# 0.02s/it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
